import pybullet as p
import time
import pybullet_data
import math
import numpy as np

# Connect to physics engine
p.connect(p.GUI)
p.setGravity(0,0,0)
p.setAdditionalSearchPath(pybullet_data.getDataPath())

# Start simulation
p.setRealTimeSimulation(0)

# Set the physics engine parameter for fixed timestep (e.g., 1/240 for slower simulation)
p.setPhysicsEngineParameter(fixedTimeStep=1/240)

# Load a built-in ground plane
plane_id = p.loadURDF("plane.urdf")


# p.configureDebugVisualizer(p.COV_ENABLE_SINGLE_STEP_RENDERING)
p.resetDebugVisualizerCamera(cameraDistance=0.2,cameraYaw=0,\
                             cameraPitch=-40,cameraTargetPosition=[0.0,-0.5,0.1])


# Load the URDF file for the robot
start_pos = [0, 0, 0.3]
start_orientation = p.getQuaternionFromEuler([0, 0, 0])
robot_id = p.loadURDF("dingo.urdf", start_pos, start_orientation)


# Set the target position for the joint
#target_position = 1.0  # Target position to move towards
#position_gain = 1  # Position control gain for smooth movement
velocity_gain = 1.0  # Velocity control gain
max_velocity = 1.0  # Set the maximum velocity (rad/s)
#max_force = 10  # Maximum force applied by the motor
# 腿的尺寸
L1 = 0.13  # 大腿长度
L2 = 0.138  # 小腿长度
# 步态参数
step_period = 1.0  # 步态周期
amplitude_x = 0.01  # 水平摆动幅度, 0.1
amplitude_z = 0.05  # 垂直摆动幅度, 
frequency = 2 * math.pi / step_period  # 频率



# Define control functions
def set_joint_angle(robot_id, joint_index, angle,position_gain):
    p.setJointMotorControl2(
        bodyUniqueId=robot_id,
        jointIndex=joint_index,
        controlMode=p.POSITION_CONTROL,
        targetPosition=angle,
        positionGain=position_gain,
        velocityGain=velocity_gain,
        #maxVelocity=max_velocity,
        #force=max_force
    )

def stand(robot_id):
    # Perform a simulation loop
    num_steps = 60  # Number of simulation steps to run
    for _ in range(num_steps): 
        set_joint_angle(robot_id, 1, 0.3, 0.1)  # 0.3/0.2 make it stand stable
        set_joint_angle(robot_id, 2, 0.3, 0.1)
        set_joint_angle(robot_id, 5, 0.3, 0.1)
        set_joint_angle(robot_id, 6, 0.3, 0.1)
        set_joint_angle(robot_id, 9, 0.3, 0.1)
        set_joint_angle(robot_id, 10, 0.3, 0.1)
        set_joint_angle(robot_id, 13, 0.3, 0.1)
        set_joint_angle(robot_id, 14, 0.3, 0.1)
        p.stepSimulation()
        time.sleep(1./240.)
        #print("I am standing...")
       
 
def all_zero(robot_id):
    # Set each joint to standing position
    for i in range(p.getNumJoints(robot_id)):
        set_joint_angle(robot_id, i, 0.0, 1)
        #print("Joint id: ", i) 

def swing_leg(robot_id):
    # Set each joint to standing position
    # for i in range(p.getNumJoints(robot_id)):
    #     if i == 5:
    #         set_joint_angle(robot_id, i, -1, 1)
    #     elif i == 6:    
    #         set_joint_angle(robot_id, i, 1, 1)
    #     else:
    #         set_joint_angle(robot_id, i, 0, 1)    
    set_joint_angle(robot_id, 5, -0.5, 1)
    set_joint_angle(robot_id, 6, 0.5, 1)

 # March in place, on going...
def March(robot_id):
        num_steps = 60  # Number of simulation steps to run
        for _ in range(num_steps):
            set_joint_angle(robot_id, 1, -0.5,0.1)  
            set_joint_angle(robot_id, 13, -0.5,0.1) 
            set_joint_angle(robot_id, 2, 0.5,0.1)  
            set_joint_angle(robot_id, 14, 0.5,0.1)  
            set_joint_angle(robot_id, 5, 0,0.1)
            set_joint_angle(robot_id, 6, 0.5,0.1) 
            set_joint_angle(robot_id, 9, 0,0.1)
            set_joint_angle(robot_id, 10, 0.5,0.1)
            p.stepSimulation()
            time.sleep(1./240.)
        num_steps = 60  # Number of simulation steps to run
        for _ in range(num_steps):    
            set_joint_angle(robot_id, 1, 0,0.1)  
            set_joint_angle(robot_id, 13, 0,0.1) 
            set_joint_angle(robot_id, 2, 0.5,0.1)  
            set_joint_angle(robot_id, 14, 0.5,0.1)  
            set_joint_angle(robot_id, 5, -0.5,0.1)
            set_joint_angle(robot_id, 6, 0.5,0.1) 
            set_joint_angle(robot_id, 9, -0.5,0.1)
            set_joint_angle(robot_id, 10, 0.5,0.1)
            p.stepSimulation() 
            time.sleep(1./240.)      
               
def lying_down(robot_id):
    # Set each joint to standing position
    for i in range(p.getNumJoints(robot_id)):
        #set knee parallel to leg
        if i == 2 or 6 or 8 or 11:
            set_joint_angle(robot_id, i, 0.8, 1)
        else:    
            set_joint_angle(robot_id, i, 0, 1)
       # print("Joint id: ", i)        
        
def walk(robot_id):
    # Example walking cycle
    for _ in range(5):  # Repeat the cycle 5 times
        ##swig leg
        for angle in [0.5, 0]:
            num_steps = 120  # Number of simulation steps to run
            for _ in range(num_steps):
                #swing leg
                set_joint_angle(robot_id, 1, angle, 0.1)
                set_joint_angle(robot_id, 2, angle, 0.1)
                set_joint_angle(robot_id, 13, angle, 0.1)
                set_joint_angle(robot_id, 14, angle, 0.1)
                #support leg
                set_joint_angle(robot_id, 5, 0.5, 0.1)
                set_joint_angle(robot_id, 6, 0.5, 0.1)
                set_joint_angle(robot_id, 9, 0.5, 0.1)
                set_joint_angle(robot_id, 10, 0.5, 0.1)
                p.stepSimulation()
                time.sleep(1./240.)
        ##step 1, 2.
        ##reverse
        for angle in [0.5, 0]:
            num_steps = 120  # Number of simulation steps to run
            for _ in range(num_steps):
                #swing leg
                set_joint_angle(robot_id, 1, 0.5, 0.1)
                set_joint_angle(robot_id, 2, 0.5, 0.1)
                set_joint_angle(robot_id, 13, 0.5, 0.1)
                set_joint_angle(robot_id, 14, 0.5, 0.1)
                #support leg
                set_joint_angle(robot_id, 5, angle, 0.1)
                set_joint_angle(robot_id, 6, angle, 0.1)
                set_joint_angle(robot_id, 9, angle, 0.1)
                set_joint_angle(robot_id, 10, angle, 0.1)
                p.stepSimulation()
                time.sleep(1./240.)
### 对角 trait ##################
# 逆运动计算, hip motor, soft 0 is parallel to body. Knee motor vertical to hip
def inverse_kinematics(x, z,leg_joint_id):
    """  
    已知足端 x ,z 坐标和 leg joint id, 计算theta_hip, theta_knee 

    参数:  
   (x, z): 足端相对于大腿关节的位置 )  
    leg joint id  
    

    返回:  
    theta_hip, theta_knee (software value,should within 0~ PI )
    """  
    # Get foot postion (x,Z) related to hip joint , 7 is FR foot @ stand mode. 5 is FR Hip motor
    #for real robot, init stand mode decide joint position, 可从正运动学推导, both start and joint pos need update
    #convert x, z axis to x, z distance to leg joint
    leg_x, leg_z = get_joint_position(robot_id, leg_joint_id)
    x = x-leg_x
    z = z-leg_z
    # print ("leg joint X axis in inverse kine, id : ", leg_x,leg_joint_id)
    # print ("leg joint z axis in inverse kine, id : ", leg_z,leg_joint_id)
    # print ("X distance: ", x)
    # print ("Z distance : ", z)
    
    D = math.sqrt(x**2 + z**2)
    theta_knee = math.acos((-D**2 + L1**2 + L2**2) / (2 * L1 * L2))
    # print ("D: ", D)
    # print ("theta_knee: ", theta_knee)
    # convert from act angle to software angle, so:
    theta_knee = math.pi/2 - theta_knee 
    # print ("theta_knee after convert Pi/2 - knee: ", theta_knee) 
    # math.atan2() get angle according to 象限 (0~Pi), (0~-pi)
    #update alpha value here...keep alpha (0,90)
    alpha = math.atan2(abs(z), abs(x))
    beta = math.acos((L1**2 + D**2 - L2**2) / (2 * L1 * D))
    # print ("alpha: ", alpha)
    # print ("beta: ", beta)
    # 0 is axis of leg joint, < 0 means leg at left side of joint, x here is (x-leg_x)
    if x < 0:
        theta_hip = alpha - beta
    else:
        theta_hip = math.pi - alpha - beta    
    return theta_hip, theta_knee


def forward_kinematics(theta_hip, theta_knee):  
    """  
    已知alpha,beta, 计算足端位置 (x, z) 相对于大腿关节的相对距离 

    参数:  
    theta1: 大腿关节角度 (度数)  
    theta2: 小腿关节角度 (度数)  
    L1: 大腿长度  
    L2: 小腿长度  

    返回:  
    (x, z): 足端相对于大腿关节的位置  
    """  

    # 将角度转换为弧度  
    # theta1_rad = math.radians(theta1)  
    # theta2_rad = math.radians(theta2)  

    # 计算足端位置 (x, z)，if/else或许可以统一 ， < 90, x<0, else x>0 , Z always > 0
    # convert from software angle to actual angle, there's delta on knee angle, software 0 is Pi/2
    theta_knee =  math.pi/2 - theta_knee 
    if theta_hip < math.pi/2:  
        x = -abs(L1 * math.cos(theta_hip) - L2 * math.cos(theta_knee - theta_hip))  
        z = L1 * abs(math.sin(theta_hip)) + L2 * abs(math.sin(theta_knee - theta_hip))  
    else:
        x = abs(L1 * math.cos(math.pi- theta_hip) + L2 * math.cos(theta_hip - theta_knee))
        ###根据几何作图，theta_hip > theta_knee -, else + , so use - 归一 
        z = abs(L1* math.sin(math.pi - theta_hip ) - L2*math.sin(theta_hip - theta_knee))
    return x, z  



def walk_diagnal(robot_id):
    #move swing leg
    #update start point
    #move support leg
    #repeat: move swing leg
    
    # step 1, decide inti postion ,eg stand mode
    # step 2, decide walk trait, eg diagnal, get x,z axis of leg joint by cal for real robot
    # step 3, get  stand mode, init position foot start axis of FR,RL  
    # step 4, generate traj. Assume leg joint fixed in one step walking
    #step 1
    stand(robot_id) 
    
    # 步态参数
    # step_period = 1.0  # 步态周期
    # amplitude_x = 0.6  # 水平摆动幅度, 0.1
    # amplitude_y = 0.29  # 垂直摆动幅度, 
    # frequency = 2 * math.pi / step_period  # 频率
    # t = time.time()  # 获取当前时间
    # print("t:  ", t)
    # 初始点和最终点，都在地上（z=0）  
    # x_start_point_FR , z_start_point_FR = get_joint_position(robot_id, 7)
    # x_start_point_RL , z_start_point_RL = get_joint_position(robot_id, 11)
    step_size = 0.05  # 步长
    num_steps = 50  # 轨迹的步数 
    
    
    #step 2,assume 4 legs are same, foot joint distance to leg joint
    x_dis,z_dis = forward_kinematics(0.3,0.3)
    #print(f"FR Foot Position by forward kin:", x_dis, z_dis)
    #step 3, (0,0,0) is floor body center , redo by cal.
    x_FR_leg, z_FR_leg = get_joint_position(robot_id, 5)
    print(f"x_FR_leg, z_FR_leg :", x_FR_leg, z_FR_leg)
    start_point_FR = x_FR_leg + x_dis
    print(f"x_FR_leg start point :", start_point_FR)
    x_RL_leg, z_RL_leg = get_joint_position(robot_id, 9)
    start_point_RL = x_RL_leg + x_dis   
    
    #step 4, 生成front right, back left leg足端轨迹
    #generate FR trajectory, offline, init stand mode, diag walk trait
    x_traj_FR, z_traj_FR = generate_foot_trajectory(start_point_FR,start_point_FR+step_size, num_steps)
    joint_angles_FR = generate_traj_angles(x_traj_FR,z_traj_FR,5)
    #update start point
    #start_point_FR = start_point_FR + step_size
    print(f"X traj :", x_traj_FR)
    print(f"z traj :", z_traj_FR)
    print(f"joint angles  :", joint_angles_FR)

    #generate RL trajectory, offline, may cal from FR data later
    x_traj_RL, z_traj_RL = generate_foot_trajectory(start_point_RL, start_point_RL+step_size, num_steps)  
    joint_angles_RL = generate_traj_angles(x_traj_RL,z_traj_RL,9)
    #update start point
    #start_point_RL = start_point_RL + step_size
        
    #move swing leg
    for i, ((theta_hip_fr, theta_knee_fr), (theta_hip_rl, theta_knee_rl)) in enumerate(zip(joint_angles_FR, joint_angles_RL)):      
        #print ("Total 100 steps, step i: ", i)
        # print ("theta hip: ",theta_hip_fl )
        # print ("theta knee: ",theta_knee_fl )
        num_steps = 50  # Number of simulation steps to run
        for _ in range(num_steps):
            # 设置伺服电机位置
            set_joint_angle(robot_id, 5, theta_hip_fr, 0.1)
            set_joint_angle(robot_id, 6, theta_knee_fr, 0.1)
            set_joint_angle(robot_id, 9, theta_hip_rl, 0.1)
            set_joint_angle(robot_id, 10, theta_knee_rl, 0.1)
            p.stepSimulation()
            time.sleep(1./240.) 



# Function to check if feet are touching the floor
def check_feet_on_floor(robot_id, plane_id):
    # Define the link indices of the feet
    foot_links = [3, 7, 11, 15]  # Adjust these to the correct foot link indices for your robot
    contact_points = p.getContactPoints(bodyA=robot_id, bodyB=plane_id)
    
    feet_on_floor = {link: False for link in foot_links}
    
    for contact in contact_points:
        link_index = contact[3]
        if link_index in foot_links:
            feet_on_floor[link_index] = True
    
    return all(feet_on_floor.values())

# Check if all feet are on the floor
if check_feet_on_floor(robot_id, plane_id):
    print("All feet are touching the floor.")
else:
    print("Not all feet are touching the floor.")
########################## end check ######################################

def generate_foot_trajectory(start_x, start_z, end, num_steps):  
    """  
    输入x 起点终点坐标, no. of steps, 生成一条腿足端轨迹, offline  
    """  
    x0 = start_x
    z0 = start_z  
    xf = end  
    
    # 线性插值生成中间点  
    x_trajectory = np.linspace(x0, xf, num_steps)  
    z_trajectory = amplitude_z * np.sin(np.pi * (x_trajectory - x0) / (xf - x0)*1.2) + z0    
                
    return x_trajectory, z_trajectory

def generate_traj_angles(x_trajectory,z_trajectory,leg_joint_id):

    joint_angles = []  
    for x, z in zip(x_trajectory, z_trajectory):  
        theta1, theta2 = inverse_kinematics(x, z,leg_joint_id)  
        joint_angles.append((theta1, theta2)) 
        
    return joint_angles
    
    
# Function to get the x, y coordinates of a specific joint (e.g., foot)  
def get_joint_position(robot_id, joint_index):  
    # Get the current state of the robot  
    state = p.getLinkState(robot_id, joint_index, computeLinkVelocity=1)  

    # Extract the position which contains (x, y, z) coordinates  
    position = state[0]  

    # Extract and return the (x, y) coordinates  
    x = position[0]  
    y = position[1]  
    z = position[2]
    return x, z  

def get_leg_length():
    #set gravity to 0, comment stand func to set all joint in software 0 states
    x_FR_leg, z_FR_leg = get_joint_position(robot_id, 5)
    x_FR, z_FR = get_joint_position(robot_id, 7)
    print(f"FR leg joint Position :", x_FR_leg, z_FR_leg)
    print(f"FR Foot Position :", x_FR, z_FR)
   
def get_joint_angle(robot_id,joint_index):
    joint_state = p.getJointState(robot_id, joint_index)  
    joint_angle = joint_state[0]  # jointState[0] is the joint angle  
    print(f"Joint {joint_index}: Angle = {joint_angle:.2f} radians")  

# my update:
# 1) generate traj, adding z para

step_size = 0.05  # 步长
num_steps = 50  # 轨迹的步数 

#step 1
stand(robot_id) 
# num_steps = 80  # Number of simulation steps to run
# for _ in range(num_steps):
#     set_joint_angle(robot_id, 5, 2, 0.1) 
#     set_joint_angle(robot_id, 6, 0.5, 0.1) 
#     p.stepSimulation()
# 初始点和最终点，都在地上（z=0）, when simulation, z > 0  
x_start_point_FR , z_start_point_FR = get_joint_position(robot_id, 7)
x_start_point_RL , z_start_point_RL = get_joint_position(robot_id, 11)
print(f"init foot position x_FR_foot, Z_FR_foot :", x_start_point_FR, z_start_point_FR)
###for test:
theta1, theta2 = inverse_kinematics(x_start_point_FR, z_start_point_FR,5)
print(f"theta1, theta2, expect 0.3 :", theta1, theta2)
get_joint_angle(robot_id,5)
get_joint_angle(robot_id,6)

#step 2,assume 4 legs are same, foot joint distance to leg joint
x_dis,z_dis = forward_kinematics(0.3,0.3)
print(f"Foot to leg distance by forward kine:", x_dis, z_dis)
#step 3, (0,0,0) is floor body center , redo by cal.
x_FR_leg, z_FR_leg = get_joint_position(robot_id, 5)
print(f"x_FR_leg, z_FR_leg :", x_FR_leg, z_FR_leg)

### 2 value conflict, use get value by temp
# start_point_FR = x_FR_leg + x_dis
# print(f"x_start_point_FR by cal :", start_point_FR)
# x_start_point_FR , z_start_point_FR = get_joint_position(robot_id, 7)
# print(f"x_start_point_FR by get :", x_start_point_FR)
# x_RL_leg, z_RL_leg = get_joint_position(robot_id, 9)
# start_point_RL = x_RL_leg + x_dis



#step 4
#convert x, z axis to x, z distance to leg joint of each legs

#generate FR trajectory, offline, init stand mode, diag walk trait
x_traj_FR, z_traj_FR = generate_foot_trajectory(x_start_point_FR,z_start_point_FR,x_start_point_FR+step_size, num_steps)
joint_angles_FR = generate_traj_angles(x_traj_FR,z_traj_FR,5)




print(f"X traj :", x_traj_FR)
print(f"z traj :", z_traj_FR)
# print(f"joint angles  :", joint_angles_FR)

#generate RL trajectory, offline
x_traj_RL, z_traj_RL = generate_foot_trajectory(x_start_point_RL, z_start_point_RL, x_start_point_RL+step_size, num_steps)  
joint_angles_RL = generate_traj_angles(x_traj_RL,z_traj_RL,9)

# Simulation loop
try:
    while True:
        # if not dog_standing:
        #     stand(robot_id)
        #     dog_standing = True
        #     print("I am standing.......")
            
          
                
        
        
        #walk_diagnal(robot_id)  
        
        
              
        
        
        #move swing leg
        # for i, ((theta_hip_fr, theta_knee_fr), (theta_hip_rl, theta_knee_rl)) in enumerate(zip(joint_angles_FR, joint_angles_RL)):      
        #     print ("Total 100 steps, step i: ", i)
        #     print ("theta hip: ",theta_hip_fr )
        #     print ("theta knee: ",theta_knee_fr )
        #     num_steps = 100  # Number of simulation steps to run
        #     for _ in range(num_steps):
        #         # 设置伺服电机位置
        #         set_joint_angle(robot_id, 5, theta_hip_fr, 0.1)
        #         set_joint_angle(robot_id, 6, theta_knee_fr, 0.1)
        #         # set_joint_angle(robot_id, 9, theta_hip_rl, 0.1)
        #         # set_joint_angle(robot_id, 10, theta_knee_rl, 0.1)
        #         p.stepSimulation()
        #         time.sleep(1./240.) 
             
         time.sleep(2)     
             
             
                
except KeyboardInterrupt:
    # Disconnect from the physics engine
    p.disconnect()
    
    
## TODO & issue/Fix
## traj seems correct, but from simulator foot is not on GND, may reverse kin issue    
# Program: invert kinect issue, use 0.078 to judge ( optimized), use foot_traj_test.py to debug
# Fix: invert kinect issue, x,y value are diff value to leg joint, not abs value.
# inverse kine issue show again, init stand mode get 0.9/0.3 
#problem 1, all joint set to 0.3 make IK fun() not correct
#Problem 2, IK abs() func issue.
#problem 3, need run step sim in loop, otherwise joint will not move in desire angle.
#solved. but see max 0.1 rad error for IK calculation.

# programming interput into stack....
# inverse kine issue show again, init stand mode get 0.9/0.3 
# inveser kinec abs issue. solved. but see max 0.1 rad error for IK calculation.
#